import warnings

warnings.filterwarnings("ignore")

from ..models.vit import VisionTransformer, interpolate_pos_embed
from ..models.med import BertConfig, BertModel, BertLMHeadModel
from transformers import BertTokenizer

import torch
from torch import nn

import os
from urllib.parse import urlparse
from timm.models.hub import download_cached_file


class BLIP_Base(nn.Module):
    def __init__(self,
                 med_config='../configs/med_config.json',
                 image_size=224,
                 vit='base',  # vision transformer 的类型
                 vit_grad_ckpt=False,
                 vit_ckpt_layer=0,
                 ):
        """
        Args:
            med_config (str):  encoder-decoder 混合模型的配置文件路径
            image_size (int): 输入图像的尺寸
            vit (str): vision transformer 的模型大小 (base 或 large)
        """
        super().__init__()

        # 创建 visual encoder (ViT)
        self.visual_encoder, vision_width = create_vit(vit, image_size, vit_grad_ckpt, vit_ckpt_layer)
        # 初始化 tokenizer
        self.tokenizer = init_tokenizer()
        # 加载 Bert 的配置
        med_config = BertConfig.from_json_file(med_config)
        # 设置 encoder 的宽度 (与 vision transformer 的输出维度一致)
        med_config.encoder_width = vision_width
        # 创建 text encoder (BertModel)
        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)

    def forward(self, image, caption, mode):

        assert mode in ['image', 'text', 'multimodal'], "mode 参数必须是 image, text, 或 multimodal"
        # 将文本转换为 token IDs
        text = self.tokenizer(caption, return_tensors="pt")  # .to(image.device)

        if mode == 'image':
            # 返回图像特征
            image_embeds = self.visual_encoder(image)
            return image_embeds

        elif mode == 'text':
            # 返回文本特征
            text_output = self.text_encoder(text.input_ids, attention_mask=text.attention_mask,
                                            return_dict=True, mode='text')
            return text_output.last_hidden_state

        elif mode == 'multimodal':
            # 返回多模态特征
            image_embeds = self.visual_encoder(image)
            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long)  # .to(image.device)

            # 将文本输入的第一个 token ID 设置为 [ENC] token ID
            text.input_ids[:, 0] = self.tokenizer.enc_token_id
            # 使用 text encoder 处理文本，并结合图像特征
            output = self.text_encoder(text.input_ids,
                                       attention_mask=text.attention_mask,
                                       encoder_hidden_states=image_embeds,  # 图像特征作为 encoder 的输入
                                       encoder_attention_mask=image_atts,  # 图像特征的 attention mask
                                       return_dict=True,
                                       )
            return output.last_hidden_state


class BLIP_Decoder(nn.Module):
    def __init__(self,
                 med_config='../configs/med_config.json',  # encoder-decoder 混合模型的配置文件
                 image_size=384,  # 输入图像大小
                 vit='base',  # Vision Transformer 的类型 ('base' 或 'large')
                 vit_grad_ckpt=False,  # 是否使用梯度检查点
                 vit_ckpt_layer=0,  # 梯度检查点的层数
                 prompt='a picture of ',  # 提示文本
                 ):
        """
        Args:
            med_config (str): encoder-decoder 模型的配置文件路径
            image_size (int): 输入图像尺寸
            vit (str): vision transformer 的模型大小
            prompt (str): 生成文本时的起始提示
        """
        super().__init__()

        # 创建视觉编码器 (Vision Transformer)
        self.visual_encoder, vision_width = create_vit(vit, image_size, vit_grad_ckpt, vit_ckpt_layer)
        # 初始化分词器
        self.tokenizer = init_tokenizer()
        # 从配置文件加载 BERT 配置
        med_config = BertConfig.from_json_file(med_config)
        # 设置 encoder 的宽度（与 vision transformer 的输出维度一致）
        med_config.encoder_width = vision_width
        # 创建文本解码器 (BertLMHeadModel)
        self.text_decoder = BertLMHeadModel(config=med_config)

        # 初始化 prompt
        self.prompt = prompt
        # 计算 prompt 的长度 (不包括 [CLS])
        self.prompt_length = len(self.tokenizer(self.prompt).input_ids) - 1

    def forward(self, image, caption):
        """
        前向传播函数，用于计算损失

        Args:
            image: 输入图像
            caption: 对应的文本描述

        Returns:
            loss_lm: 语言建模损失
        """

        # 获取图像特征
        image_embeds = self.visual_encoder(image)
        # 创建图像 attention mask（所有位置都为 1，表示关注所有 patch）
        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image.device)

        # 对文本进行分词，并进行 padding 和截断
        text = self.tokenizer(caption, padding='longest', truncation=True, max_length=40, return_tensors="pt").to(
            image.device)

        # 将文本的第一个 token ID 设置为 [BOS] token ID
        text.input_ids[:, 0] = self.tokenizer.bos_token_id

        # 创建 decoder 的 targets, 将填充位置设为 -100 (在计算损失时会被忽略)
        decoder_targets = text.input_ids.masked_fill(text.input_ids == self.tokenizer.pad_token_id, -100)
        # prompt 部分的 targets 也设置为 -100 (不计算 prompt 部分的损失)
        decoder_targets[:, :self.prompt_length] = -100

        # 使用 text decoder 进行解码
        decoder_output = self.text_decoder(text.input_ids,
                                           attention_mask=text.attention_mask,
                                           encoder_hidden_states=image_embeds,  # 图像特征作为 encoder 的输入
                                           encoder_attention_mask=image_atts,  # 图像的 attention mask
                                           labels=decoder_targets,  # decoder 的 targets
                                           return_dict=True,
                                           )
        # 获取语言建模损失
        loss_lm = decoder_output.loss

        return loss_lm

    def generate(self, image, sample=False, num_beams=3, max_length=30, min_length=10, top_p=0.9,
                 repetition_penalty=1.0):
        """
        生成文本描述

        Args:
            image: 输入图像
            sample: 是否使用采样 (True) 或束搜索 (False)
            num_beams: 束搜索的束数量
            max_length: 生成文本的最大长度
            min_length: 生成文本的最小长度
            top_p: nucleus sampling 的 top-p 值
            repetition_penalty: 重复惩罚

        Returns:
            captions: 生成的文本描述列表
        """
        # 获取图像 embedding
        image_embeds = self.visual_encoder(image)

        if not sample:
            # 如果不使用采样，则将图像 embedding 复制 num_beams 份 (用于束搜索)
            image_embeds = image_embeds.repeat_interleave(num_beams, dim=0)

        # 创建图像 attention mask
        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image.device)
        # 创建模型参数字典
        model_kwargs = {"encoder_hidden_states": image_embeds, "encoder_attention_mask": image_atts}

        # 将 prompt 复制多份 (与输入图像数量相同)
        prompt = [self.prompt] * image.size(0)
        # 对 prompt 进行分词，并获取 input_ids
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(image.device)
        # 将 prompt 的第一个 token ID 设置为 [BOS] token ID
        input_ids[:, 0] = self.tokenizer.bos_token_id
        # 去掉 prompt 的最后一个 token (通常是 [SEP])
        input_ids = input_ids[:, :-1]

        if sample:
            # 使用 nucleus sampling 进行生成
            outputs = self.text_decoder.generate(input_ids=input_ids,
                                                 max_length=max_length,
                                                 min_length=min_length,
                                                 do_sample=True,
                                                 top_p=top_p,
                                                 num_return_sequences=1,
                                                 eos_token_id=self.tokenizer.sep_token_id,  # 设置结束 token
                                                 pad_token_id=self.tokenizer.pad_token_id,  # 设置填充 token
                                                 repetition_penalty=1.1,  # 设置重复惩罚
                                                 **model_kwargs)
        else:
            # 使用 beam search 进行生成
            outputs = self.text_decoder.generate(input_ids=input_ids,
                                                 max_length=max_length,
                                                 min_length=min_length,
                                                 num_beams=num_beams,
                                                 eos_token_id=self.tokenizer.sep_token_id,
                                                 pad_token_id=self.tokenizer.pad_token_id,
                                                 repetition_penalty=repetition_penalty,
                                                 **model_kwargs)

        # 从 output 中提取生成的文本
        captions = []
        for output in outputs:
            caption = self.tokenizer.decode(output, skip_special_tokens=True)  # 解码
            captions.append(caption[len(self.prompt):])  # 去掉 prompt 部分
        return captions


def blip_decoder(pretrained='', **kwargs):
    """
    创建并加载 BLIP Decoder 模型

    Args:
        pretrained (str): 预训练模型的路径或 URL
        **kwargs: 其他模型参数

    Returns:
        model: BLIP Decoder 模型
    """
    model = BLIP_Decoder(**kwargs)
    if pretrained:
        # 如果提供了预训练模型路径，则加载预训练权重
        model, msg = load_checkpoint(model, pretrained)
        assert (len(msg.missing_keys) == 0)
    return model


def blip_feature_extractor(pretrained='', **kwargs):
    """
     创建并加载 BLIP 特征提取器（Base）模型。

     Args:
         pretrained (str): 预训练模型的路径或 URL。
         **kwargs: 其他模型参数。

     Returns:
         model: BLIP_Base 模型。
    """
    model = BLIP_Base(**kwargs)
    if pretrained:
        model, msg = load_checkpoint(model, pretrained)
        assert (len(msg.missing_keys) == 0)
    return model


def init_tokenizer():
    """
    初始化分词器

    Returns:
        tokenizer: 分词器
    """
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    # 添加特殊 token [DEC] (bos_token)
    tokenizer.add_special_tokens({'bos_token': '[DEC]'})
    # 添加特殊 token [ENC]
    tokenizer.add_special_tokens({'additional_special_tokens': ['[ENC]']})
    # 获取 [ENC] token 的 ID
    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]
    return tokenizer


def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
    """
    创建 Vision Transformer (ViT) 模型

    Args:
        vit (str): ViT 的类型 ('base' 或 'large')
        image_size (int): 输入图像尺寸
        use_grad_checkpointing (bool): 是否使用梯度检查点
        ckpt_layer (int): 梯度检查点的层数
        drop_path_rate (float): drop path 的概率

    Returns:
        visual_encoder: ViT 模型
        vision_width: ViT 输出的特征维度
    """

    assert vit in ['base', 'large'], "vit 参数必须是 base 或 large"
    if vit == 'base':
        vision_width = 768
        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12,
                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing,
                                           ckpt_layer=ckpt_layer,
                                           drop_path_rate=0 or drop_path_rate
                                           )
    elif vit == 'large':
        vision_width = 1024
        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24,
                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing,
                                           ckpt_layer=ckpt_layer,
                                           drop_path_rate=0.1 or drop_path_rate
                                           )
    return visual_encoder, vision_width


def is_url(url_or_filename):
    """
    判断给定的字符串是否是 URL

    Args:
        url_or_filename (str): 要判断的字符串

    Returns:
        bool: 如果是 URL，则返回 True，否则返回 False
    """
    parsed = urlparse(url_or_filename)
    return parsed.scheme in ("http", "https")


def load_checkpoint(model, url_or_filename):
    """
      加载预训练模型的权重。

      Args:
          model: 要加载权重的模型。
          url_or_filename: 预训练模型的 URL 或本地文件路径。

      Returns:
          model: 加载了权重的模型。
          msg:  load_state_dict 的返回值 (包含 missing_keys 和 unexpected_keys)。
    """
    if is_url(url_or_filename):
        # 如果是 URL，则下载预训练模型
        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
        checkpoint = torch.load(cached_file, map_location='cpu')
    elif os.path.isfile(url_or_filename):
        # 如果是本地文件，则直接加载
        checkpoint = torch.load(url_or_filename, map_location='cpu')
    else:
        raise RuntimeError('checkpoint url 或路径无效')

    state_dict = checkpoint['model']

    # 对视觉编码器的位置嵌入进行插值
    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],
                                                                   model.visual_encoder)
    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
                                                                         model.visual_encoder_m)
    # 删除形状不匹配的权重
    for key in model.state_dict().keys():
        if key in state_dict.keys():
            if state_dict[key].shape != model.state_dict()[key].shape:
                del state_dict[key]

    msg = model.load_state_dict(state_dict, strict=False)
    print('从 %s 加载 checkpoint' % url_or_filename)
    return model, msg
